import heapq

def move(loc, dir):
    # 根据方向移动位置
    directions = [
        (0, 0),  # 原地不动
        (0, -1),  # 上
        (1, 0),  # 右
        (0, 1),  # 下
        (-1, 0),  # 左
    ]

    return loc[0] + directions[dir][0], loc[1] + directions[dir][1]


def get_sum_of_cost(paths):
    # 计算所有路径的总成本
    rst = 0
    for path in paths:
        rst += len(path) - 1
    return rst


def compute_heuristics(my_map, goal):
    # 使用Dijkstra算法从目标位置构建最短路径树
    open_list = []
    closed_list = dict()
    root = {'loc': goal, 'cost': 0}
    heapq.heappush(open_list, (root['cost'], goal, root))
    closed_list[goal] = root
    while len(open_list) > 0:
        (cost, loc, curr) = heapq.heappop(open_list)
        for dir in range(5):
            child_loc = move(loc, dir)
            child_cost = cost + 1
            if child_loc[0] < 0 or child_loc[0] >= len(my_map) \
               or child_loc[1] < 0 or child_loc[1] >= len(my_map[0]):
               continue
            if my_map[child_loc[0]][child_loc[1]]:
                continue
            child = {'loc': child_loc, 'cost': child_cost}
            if child_loc in closed_list:
                existing_node = closed_list[child_loc]
                if existing_node['cost'] > child_cost:
                    closed_list[child_loc] = child
                    # open_list.delete((existing_node['cost'], existing_node['loc'], existing_node))
                    heapq.heappush(open_list, (child_cost, child_loc, child))
            else:
                closed_list[child_loc] = child
                heapq.heappush(open_list, (child_cost, child_loc, child))

    # 构建启发式值表
    h_values = dict()
    for loc, node in closed_list.items():
        h_values[loc] = node['cost']
    return h_values


def build_constraint_table(constraints, agent):
    # 返回一个包含给定代理人的约束列表的表格，按时间步索引
    table = {}
    length = len(constraints)

    for i in range(length):
        if constraints[i]['agent'] == agent:

            if constraints[i]['timestep'] in table:
                # table[constraints[i]['timestep']].append(constraints[i]['loc'])
                table[constraints[i]['timestep']].append({'locs': constraints[i]['loc'], 'positive': constraints[i]['positive']})

            else:
                # table[constraints[i]['timestep']] = [constraints[i]['loc']]
                table[constraints[i]['timestep']] = [{'locs': constraints[i]['loc'], 'positive': constraints[i]['positive']}]

    return table


def get_location(path, time):
    # 获取路径中指定时间步的位置
    if time < 0:
        return path[0]
    elif time < len(path):
        return path[time]
    else:
        return path[-1]  # wait at the goal location


def get_path(goal_node):
    # 从目标节点回溯获取路径
    path = []
    curr = goal_node
    while curr is not None:
        path.append(curr['loc'])
        curr = curr['parent']
    path.reverse()
    return path


def is_constrained(curr_loc, next_loc, next_time, constraint_table):
    # 检查从curr_loc到next_loc在时间步next_time的移动是否违反任何约束
    constraints_list = constraint_table.get(next_time)

    if not constraints_list:
        return False, False

    ret = False
    positive = False
    for i in range(len(constraints_list)):

        if constraints_list[i]['positive'] == True:
            positive = True
        else:
            positive = False

        # check for vertex contraint
        if constraint_table and constraints_list[i] and (next_loc in constraints_list[i]['locs']) and len(constraints_list[i]['locs']) == 1:
            ret = True

        # check for edge constraint
        elif constraint_table and constraints_list[i] and (curr_loc in constraints_list[i]['locs']) and (next_loc in constraints_list[i]['locs']):
            ret = True

        # no constraints found
        else:
            ret = False

        if ret:
            return True, positive

    return False, positive


def push_node(open_list, node):
    # 将节点推入开放列表
    heapq.heappush(open_list, (node['g_val'] + node['h_val'], node['h_val'], node['loc'], node))


def pop_node(open_list):
    # 从开放列表中弹出节点
    _, _, _, curr = heapq.heappop(open_list)
    return curr


def compare_nodes(n1, n2):
    """Return true is n1 is better than n2."""
    return n1['g_val'] + n1['h_val'] < n2['g_val'] + n2['h_val']


def a_star(my_map, start_loc, goal_loc, h_values, agent, constraints):
    """ my_map      - binary obstacle map
        start_loc   - start position
        goal_loc    - goal position
        agent       - the agent that is being re-planned
        constraints - constraints defining where robot should or cannot go at each timestep
    """
    # 初步检查起点和目标点的有效性
    if (my_map[start_loc[0]][start_loc[1]] or 
        my_map[goal_loc[0]][goal_loc[1]]):
        print(f"Agent {agent}: 起点或目标点在障碍物上")
        return None

    # 检查起点到目标点的可达性
    if start_loc not in h_values or goal_loc not in h_values:
        print(f"Agent {agent}: 目标点 {goal_loc} 无法从起点 {start_loc} 到达")
        return None

    # 新增：检查起点是否可达目标
    if start_loc not in h_values:
        print(f"Agent {agent}: 目标点 {goal_loc} 无法从起点 {start_loc} 到达（被障碍物阻挡）")
        return None

    # calculate upper bound on path length for an agent
    environment_size = 0
    for list in my_map:
        environment_size += list.count(0)

    # path_length_upper_bound = 13

    # 在空间-时间域中搜索，而不是仅在空间域中
    biggest_constraint_ts = 0
    for c in constraints:
        if c['agent'] == agent and c['timestep'] > biggest_constraint_ts:
            biggest_constraint_ts = c['timestep']

    earliest_goal_timestep = biggest_constraint_ts

    # 构建约束表
    constraint_table = build_constraint_table(constraints, agent)

    open_list = []
    closed_list = dict()
    h_value = h_values[start_loc]
    root = { 'loc': start_loc, 'g_val': 0, 'h_val': h_value, 'parent': None, 'ts': 0 }
    push_node(open_list, root)
    closed_list[(root['loc'], root['ts'])] = root

    while len(open_list) > 0:
        curr = pop_node(open_list)

        # exceeded path length upper bound
        # if len(get_path(curr)) >= path_length_upper_bound:
            # return None

        # 调整目标测试条件以处理目标约束
        if (curr['loc'] == goal_loc) and (curr['ts'] >= earliest_goal_timestep):
            return get_path(curr)

        for dir in range(5):
            child_loc = move(curr['loc'], dir)

            # no bounds on test maps?
            if child_loc[0] > (len(my_map) - 1) or child_loc[1] > (len(my_map[0]) - 1) or child_loc[0] < 0 or child_loc[1] < 0:
                continue

            if my_map[child_loc[0]][child_loc[1]]:
                continue


            child = {'loc': child_loc,
                    'g_val': curr['g_val'] + 1,
                    'h_val': h_values[child_loc],
                    'parent': curr,
                    'ts': curr['ts'] + 1}

            # 检查新节点是否满足约束，如果不满足则剪枝
            constrained, positive = is_constrained(curr['loc'], child['loc'], child['ts'], constraint_table)

            if constrained:
                if positive == False:
                    continue

            if (child['loc'], child['ts']) in closed_list:
                existing_node = closed_list[(child['loc'], child['ts'])]
                if compare_nodes(child, existing_node):
                    closed_list[(child['loc'])] = child
                    push_node(open_list, child)
            else:
                closed_list[(child['loc'], child['ts'])] = child
                push_node(open_list, child)

    return None  # Failed to find solutions